import lightgbm as lgb
import pandas as pd
import numpy as np
import argparse
import json
import collections

def predict(modelpath,testdatapath,resultpath):
    gbm=lgb.Booster(model_file=modelpath)

    testdata=pd.read_csv(testdatapath)
    totalcount=testdata.shape[0]
    avgcount=totalcount//6

    featurelist=list(testdata.columns)
    featurelist.remove('sample_id')
    featurelist.remove('feature100')
    featurelist.remove('feature57')
    featurelist.remove('feature77')
    featurelist.remove('feature64')
    featurelist.remove('feature92')
    if('label' in featurelist):
        featurelist.remove('label')

    prediction_prob = gbm.predict(testdata[featurelist])
    prediction=np.argmax(prediction_prob,1)

    a0=prediction_prob[:,0]
    a1=prediction_prob[:,1]
    a2=prediction_prob[:,2]

    sortindex0=np.argsort(a0)[::-1]
    sortindex1=np.argsort(a1)[::-1]
    sortindex2=np.argsort(a2)[::-1]


    prediction[sortindex1[0:avgcount]]=1
    prediction[sortindex2[0:avgcount]]=2

    
    prediction=prediction.tolist()
    sampleid_list=testdata['sample_id'].values.tolist()

    resultdict=dict(zip(sampleid_list,prediction))
    resultjson=json.dumps(resultdict,ensure_ascii=False)
    # print(resultjson)
    with open(resultpath,'w',encoding='utf-8') as writejson:
        writejson.write(resultjson)

if __name__=='__main__':
    # predict('./model_2023-07-10_22-14-26.016396.txt','./data/test_2000_x.csv','result.json')
    parser = argparse.ArgumentParser()
    parser.add_argument("modelpath",default='',type=str)
    parser.add_argument("testdatapath",default='',type=str)
    parser.add_argument("resultpath",default='',type=str)
    args = parser.parse_args()
    predict(args.modelpath,args.testdatapath,args.resultpath)
